import urllib.request
from datetime import datetime
from plugins import logger, notifier
from tqdm import tqdm
import argparse


def download_file(zone, filename):
    city2url = {
        "cn": "https://metowolf.github.io/iplist/data/country/CN.txt",
        "bj": "https://metowolf.github.io/iplist/data/cncity/110000.txt"
    }
    url = city2url[zone]
    try:
        urllib.request.urlretrieve(url, filename)
        logger.info(f"已下载并保存为 {filename}")
        notifier.send_message(f"文件 {filename} 下载成功！")
    except Exception as e:
        logger.error(f"下载文件失败: {e}")
        notifier.send_message(f"文件 {filename} 下载失败，错误信息: {e}")


def convert_to_ipset_format(filename, ipset_name):
    #  实现将文件转换为ipset可以加载的格式
    print("开始转换")  # 开始转换
    with open(filename, 'r') as fi:
        with open(f"zones/{ipset_name}.ipset", 'w') as fo:
            fo.write(f"create {ipset_name} hash:net\n")
            for line in tqdm(fi.readlines()):
                new_line = f"add {ipset_name} {line}"
                fo.write(new_line)
            fo.write("\n")
            fo.write(f"add {ipset_name} 192.168.0.0/24\n")
            fo.write(f"add {ipset_name} 192.168.1.0/24\n")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='自动下载并转换IP地址库')
    parser.add_argument('-z', '--zone', required=True, help='地区代码(如: cn, bj等)')
    args = parser.parse_args()
    
    # step1: 下载北京ip地址库
    filename = f"zones/{args.zone}_{datetime.now().strftime('%Y%m%d')}.zone"
    download_file(args.zone, filename)
    # step2: 转化为ipset可以加载的格式
    ipset_name = f"{args.zone}_ips"
    convert_to_ipset_format(filename, ipset_name)

